import os
import numpy as np
import vtk
from vtkmodules.util import numpy_support
import meshio
import h5py as h5


class VTKWriter:
    def __init__(self, reader):
        """
        Initialize the VTKWriter.

        :param reader: BinReader object.

        """
        self.reader = reader

    def create_grid(self):
        """
        Create a VTK grid with points and cells.

        :return: vtk.vtkUnstructuredGrid

        """
        grid = vtk.vtkUnstructuredGrid()

        # Process points
        GRID_VTX_array = {}
        if self.reader.get_static_data("GRID_VTX", GRID_VTX_array):
            data = np.array(GRID_VTX_array["data"], dtype=np.float64)
            points_array = data.reshape(-1, 3)
            vtk_points = vtk.vtkPoints()
            vtk_points.SetData(numpy_support.numpy_to_vtk(points_array))
            grid.SetPoints(vtk_points)

        # Process cells
        NODE_LIST_array = {}
        if self.reader.get_static_data("NODE_LIST", NODE_LIST_array):
            cell_data = np.array(NODE_LIST_array["data"], dtype=np.int64)
            n_cells = len(cell_data) // 8
            cell_data = cell_data.reshape(-1, 8)
            connectivity = np.hstack(
                [np.full((n_cells, 1), 8, dtype=np.int64), cell_data]
            )
            connectivity_flat = connectivity.flatten()
            vtk_cells = vtk.vtkCellArray()
            vtk_cells.SetCells(
                n_cells, numpy_support.numpy_to_vtkIdTypeArray(connectivity_flat)
            )
            grid.SetCells(vtk.VTK_HEXAHEDRON, vtk_cells)

        return grid

    def write_vtk(self, output_folder, filename):
        """
        Write VTK files for all timesteps.

        :param output_folder: Directory path for output files.
        :param filename: Base filename for output files.
        """
        os.makedirs(output_folder, exist_ok=True)
        grid = self.create_grid()

        for ind, dynamic_infos in enumerate(self.reader.dym_props):
            tstep = int(dynamic_infos[0].time)
            if ind >= len(self.reader.dym_props):
                print(f"Time step {tstep} does not exist. Skipping.")
                continue

            grid.GetCellData().Initialize()

            for info in dynamic_infos:
                prop_array = {}
                if self.reader.get_dynamic_data(ind, info.name, prop_array):
                    if prop_array["on_grid"]:
                        if prop_array["type"] in ["d", "i"]:
                            print(
                                f"Adding property {info.name} for timestep {tstep} to grid"
                            )
                            data = numpy_support.numpy_to_vtk(
                                prop_array["data"], deep=True
                            )
                            data.SetName(info.name)
                            grid.GetCellData().AddArray(data)
                        else:
                            print(
                                f"Unknown data type {prop_array['type']} for property {info.name}"
                            )
                    else:
                        print(f"Property {info.name} is not on grid.")
                else:
                    print(f"Property {info.name} not found for timestep {tstep}.")

            timestep_filename = f"{filename}_t{tstep}.vtu"
            output_filepath = os.path.join(output_folder, timestep_filename)
            writer = vtk.vtkXMLUnstructuredGridWriter()
            writer.SetFileName(output_filepath)
            writer.SetInputData(grid)
            writer.Write()
            print(f"Successfully wrote grid for timestep {tstep} to {output_filepath}")

    def write_vtk_single(self, tstep, filename):
        """
        Write a single VTK file for a specific timestep.

        :param tstep: Time step to write.
        :param filename: Output file name.

        """
        if tstep > len(self.reader.dym_props):
            print(f"Error: Time step {tstep} does not exist.")
            return

        grid = self.create_grid()
        grid.GetCellData().Initialize()

        for info in self.reader.dym_props[tstep]:
            prop_array = {}
            if self.reader.get_dynamic_data(tstep, info.name, prop_array):
                if prop_array["on_grid"]:
                    if prop_array["type"] in ["d", "i"]:
                        print(f"Adding property {info.name} to grid")
                        data = numpy_support.numpy_to_vtk(prop_array["data"], deep=True)
                        data.SetName(info.name)
                        grid.GetCellData().AddArray(data)
                    else:
                        print(f"Unknown data type {prop_array['type']}")
                else:
                    print(f"Property {info.name} is not on grid.")
            else:
                print(f"Property {info.name} not found.")

        writer = vtk.vtkXMLUnstructuredGridWriter()
        writer.SetFileName(filename)
        writer.SetInputData(grid)
        writer.Write()
        print(f"Successfully wrote grid to {filename}")


class MeshioWriter:
    def __init__(self, reader):
        """
        Initialize the MeshioWriter.

        :param reader: BinReader object.

        """
        self.reader = reader

    def create_mesh(self):
        """
        Create a meshio Mesh with points and cells.

        :return: meshio.Mesh

        """
        # Process points
        GRID_VTX_array = {}
        if self.reader.get_static_data("GRID_VTX", GRID_VTX_array):
            data = np.array(GRID_VTX_array["data"], dtype=np.float64)
            points = data.reshape(-1, 3)
            print(f"Found {len(points)} points.")
        else:
            raise ValueError("GRID_VTX data not found.")

        # Process cells
        NODE_LIST_array = {}
        if self.reader.get_static_data("NODE_LIST", NODE_LIST_array):
            cell_data = (
                np.array(NODE_LIST_array["data"], dtype=np.int64) - 1
            )  # Convert to zero-based indices
            n_cells = len(cell_data) // 8
            cell_data = cell_data.reshape(-1, 8)
            # cell_data = cell_data[:self.reader.get_active_base_grid_num() - 1]
            # Meshio expects zero-based indices
            cells = [("hexahedron", cell_data)]

        else:
            raise ValueError("NODE_LIST data not found.")

        return meshio.Mesh(points, cells)

    def write_poro(self, output_folder, filename, file_format="xdmf"):
        """
        Write porosity data to an XDMF file. File format can be changed to VTK by changing the file_format parameter.

        :param output_folder: Directory path for output files.
        :param filename: Base filename for output files.
        :param file_format: File format for output files. Default is XDMF.
        """
        os.makedirs(output_folder, exist_ok=True)
        mesh = self.create_mesh()

        PORO_array = {}
        if self.reader.get_static_data("PORO", PORO_array):
            data = np.array(PORO_array["data"], dtype=np.float64)
            cell_data = {"porosity": [data]}
            mesh.cell_data = cell_data

            output_filepath = os.path.join(
                output_folder, f"{filename}_poro.{file_format}"
            )
            # meshio.write(output_filepath, mesh, file_format="xdmf", compression=None, data_format="HDF")
            meshio.write(output_filepath, mesh, file_format="xdmf")
            print(f"Successfully wrote porosity data to {output_filepath}")

    def write_all_timesteps_timeseries(self, output_folder, filename, file_format="xmf"):
        """
        Write all timesteps into a single XDMF file.

        :param output_folder: Directory path for output files.
        :param filename: Base filename for output file.
        :param file_format: File format for output file. Default is XDMF.
        """

        os.makedirs(output_folder, exist_ok=True)
        mesh = self.create_mesh()

        output_filepath = os.path.join(output_folder, f"{filename}.{file_format}")
        print(output_filepath)

        # Need to change meshio code: line 265, time_series.py, stem -> with_suffix
        with meshio.xdmf.TimeSeriesWriter(output_filepath, data_format="HDF") as writer:
            writer.write_points_cells(mesh.points, mesh.cells)

            for ind, dynamic_infos in enumerate(self.reader.dym_props):
                tstep = int(dynamic_infos[0].time)

                # prepare cell data
                cell_data = {}
                for info in dynamic_infos:
                    prop_array = {}
                    if self.reader.get_dynamic_data(ind, info.name, prop_array):
                        if prop_array["on_grid"]:
                            if prop_array["type"] in ["d", "i"]:
                                print(f"Adding property {info.name} for timestep {tstep} to mesh")
                                cell_data[info.name] = [prop_array["data"]]
                            else:
                                print(f"Unknown data type {prop_array['type']} for property {info.name}")
                        else:
                            print(f"Property {info.name} is not on grid.")
                    else:
                        print(f"Property {info.name} not found for timestep {tstep}.")

                # Write mesh with cell data
                writer.write_data(tstep, point_data=None, cell_data=cell_data)
                print(f"Successfully add data for timestep {tstep}")


    def write_all_timesteps(self, output_folder, filename, file_format="xdmf"):
        """
        Write XDMF files for all timesteps. File format can be changed to VTK by changing the file_format parameter.

        :param output_folder: Directory path for output files.
        :param filename: Base filename for output files.
        :param file_format: File format for output files. Default is XDMF.
        """
        os.makedirs(output_folder, exist_ok=True)
        mesh = self.create_mesh()

        for ind, dynamic_infos in enumerate(self.reader.dym_props):
            tstep = int(dynamic_infos[0].time)
            if ind >= len(self.reader.dym_props):
                print(f"Time step {tstep} does not exist. Skipping.")
                continue

            # Prepare cell data
            cell_data = {}
            for info in dynamic_infos:
                prop_array = {}
                if self.reader.get_dynamic_data(ind, info.name, prop_array):
                    if prop_array["on_grid"]:
                        if prop_array["type"] in ["d", "i"]:
                            print(
                                f"Adding property {info.name} for timestep {tstep} to mesh"
                            )
                            cell_data[info.name] = [prop_array["data"]]
                        else:
                            print(
                                f"Unknown data type {prop_array['type']} for property {info.name}"
                            )
                    else:
                        print(f"Property {info.name} is not on grid.")
                else:
                    print(f"Property {info.name} not found for timestep {tstep}.")

            timestep_filename = f"{filename}_t{tstep}.{file_format}"
            output_filepath = os.path.join(output_folder, timestep_filename)

            # Write mesh with cell data
            mesh.cell_data = cell_data
            meshio.write(output_filepath, mesh, file_format=file_format)
            print(f"Successfully wrote mesh for timestep {tstep} to {output_filepath}")

    def write_single_timestep(self, tstep, filename, file_format="xdmf"):
        """
        Write a single XDMF file for a specific timestep. File format can be changed to VTK by changing the file_format parameter.

        :param tstep: Time step to write.
        :param filename: Output file name.
        :param file_format: File format for output files. Default is XDMF.
        """
        if tstep > len(self.reader.dym_props):
            print(f"Error: Time step {tstep} does not exist.")
            return

        mesh = self.create_mesh()

        # Prepare cell data
        cell_data = {}
        dynamic_infos = self.reader.dym_props[tstep]
        for info in dynamic_infos:
            prop_array = {}
            if self.reader.get_dynamic_data(tstep, info.name, prop_array):
                if prop_array["on_grid"]:
                    if prop_array["type"] in ["d", "i"]:
                        print(f"Adding property {info.name} to mesh")
                        cell_data[info.name] = [prop_array["data"]]
                    else:
                        print(f"Unknown data type {prop_array['type']}")
                else:
                    print(f"Property {info.name} is not on grid.")
            else:
                print(f"Property {info.name} not found.")

        # Write mesh with cell data
        mesh.cell_data = cell_data
        output_filepath = f"{filename}_t{tstep}.{file_format}"
        meshio.write(output_filepath, mesh, file_format=file_format)
        print(f"Successfully wrote mesh to {output_filepath}")

    def write_well(self, output_folder, filename, file_format="xdmf"):
        """
        Write well data to an XDMF file. File format can be changed to VTK by changing the file_format parameter.

        :param output_folder: Directory path for output files.
        :param filename: Base filename for output files.
        :param file_format: File format for output files. Default is XDMF.

        """
        os.makedirs(output_folder, exist_ok=True)

        # fix relative offset
        global_xcoord = self.reader.well_data.XCoord - self.reader.AxisOffset[0]
        global_ycoord = self.reader.well_data.YCoord - self.reader.AxisOffset[1]

        prefix = np.zeros(len(self.reader.well_data.stat_data) + 1, dtype=np.int64)

        points = []
        for i, well_points in enumerate(self.reader.well_data.stat_data):
            Xcoord = well_points["Xcoord"]
            Ycoord = well_points["Ycoord"]
            Depth = well_points["Depth"]

            # add offset
            Xcoord += global_xcoord
            Ycoord += global_ycoord

            prefix[i + 1] = prefix[i] + len(Xcoord)  # prefix for each well

            well_point_data = np.stack((Xcoord, Ycoord, Depth), axis=-1)
            points.append(well_point_data)

        points = np.concatenate(points)
        cells = [("vertex", np.arange(len(points)).reshape(-1, 1))]
        mesh = meshio.Mesh(points, cells)

        for ind, dym_info in enumerate(self.reader.well_data.dym_data):
            tstep = ind
            point_data = {}
            for well_data in dym_info["Data"]:
                id = well_data["ID"] - 1
                for prop in well_data:
                    if prop == "ID":
                        continue
                    if prop not in point_data:
                        # point_data[prop] = np.zeros(self.reader.well_data.points_num, dtype=np.float64)
                        point_data[prop] = np.full(
                            self.reader.well_data.points_num, np.nan, dtype=np.float64
                        )  # fill with nan
                    point_data[prop][prefix[id] : prefix[id + 1]] = well_data[
                        prop
                    ]  # fill with well data

            timestep_filename = f"{filename}_well_t{tstep}.{file_format}"
            output_filepath = os.path.join(output_folder, timestep_filename)
            mesh.point_data = point_data
            meshio.write(output_filepath, mesh, file_format=file_format)
            print(
                f"Successfully wrote well data for timestep {tstep} to {output_filepath}"
            )


class VTKHDFWriter:
    def __init__(self, reader):
        """
        Initialize the VTKHDFWriter.

        :param reader: BinReader object.

        """
        self.reader = reader

    def create_grid(self):
        """
        Create a VTK grid with points and cells.

        :return: vtk.vtkUnstructuredGrid

        """
        grid = vtk.vtkUnstructuredGrid()
        # Process points
        GRID_VTX_array = {}
        if self.reader.get_static_data("GRID_VTX", GRID_VTX_array):
            data = np.array(GRID_VTX_array["data"], dtype=np.float64)
            points_array = data.reshape(-1, 3)
            vtk_points = vtk.vtkPoints()
            vtk_points.SetData(numpy_support.numpy_to_vtk(points_array))
            grid.SetPoints(vtk_points)

        # Process cells
        NODE_LIST_array = {}
        if self.reader.get_static_data("NODE_LIST", NODE_LIST_array):
            cell_data = np.array(NODE_LIST_array["data"], dtype=np.int32) - 1 # Convert to zero-based indices
            n_cells = len(cell_data) // 8
            cell_data = cell_data.reshape(-1, 8)
            connectivity = np.hstack(
                [np.full((n_cells, 1), 8, dtype=np.int64), cell_data]
            )
            connectivity_flat = connectivity.flatten()
            vtk_cells = vtk.vtkCellArray()
            vtk_cells.SetCells(
                n_cells, numpy_support.numpy_to_vtkIdTypeArray(connectivity_flat)
            )
            grid.SetCells(vtk.VTK_HEXAHEDRON, vtk_cells)

        return grid
    
    def write_static_grid(self, output_folder, filename):
        """
        Write VTK file for the static grid.

        :param output_folder: Directory path for output files.
        :param filename: Base filename for output files.
        """
        os.makedirs(output_folder, exist_ok=True)
        grid = self.create_grid()

        output_filepath = os.path.join(output_folder, f"{filename}_static.vtkhdf")
        writer = vtk.vtkHDFWriter()
        writer.SetFileName(output_filepath)
        writer.SetInputData(grid)
        writer.Write()
        print(f"Successfully wrote static grid to {output_filepath}")

    def write_poro(self, output_folder, filename):
        """
        Write porosity data to a vtkhdf file.

        :param output_folder: Directory path for output files.
        :param filename: Base filename for output files.
        """
        os.makedirs(output_folder, exist_ok=True)
        grid = self.create_grid()

        PORO_array = {}
        if self.reader.get_static_data("PORO", PORO_array):
            cell_data = numpy_support.numpy_to_vtk(PORO_array["data"], deep=True)
            cell_data.SetName("porosity")
            grid.GetCellData().AddArray(cell_data)

            timestep_filename = f"{filename}_poro.vtkhdf"
            output_filepath = os.path.join(output_folder, timestep_filename)
            writer = vtk.vtkHDFWriter()
            writer.SetFileName(output_filepath)
            writer.SetInputData(grid)
            writer.Write()
            print(f"Successfully wrote porosity data to {output_filepath}")

    def write_timeseries(self, output_folder, filename):
        """
        Write vtkhdf files for all timesteps.

        :param output_folder: Directory path for output files.
        :param filename: Base filename for output files.
        """
        os.makedirs(output_folder, exist_ok=True)
        grid = self.create_grid()

        timesteps = vtk.vtkGenerateTimeSteps()
        np_cell_data = {}
        nsteps = 0
        time_values = []

        for ind, dynamic_infos in enumerate(self.reader.dym_props):
            tstep = int(dynamic_infos[0].time)
            if ind >= len(self.reader.dym_props):
                print(f"Time step {tstep} does not exist. Skipping.")
                continue

            timesteps.AddTimeStepValue(tstep)
            nsteps += 1
            time_values.append(tstep)

            for info in dynamic_infos:
                prop_array = {}
                if self.reader.get_dynamic_data(ind, info.name, prop_array):
                    if prop_array["on_grid"]:
                        if prop_array["type"] in ["d", "i"]:
                            print(
                                f"Adding property {info.name} for timestep {tstep} to grid"
                            )
                            # data = numpy_support.numpy_to_vtk(
                                # prop_array["data"], deep=True
                            # )
                            # data.SetName(info.name)
                            # grid.GetCellData().AddArray(data)
                            # timesteps.AddTimeStepValue(tstep)

                            if np_cell_data.get(info.name) is None:
                                np_cell_data[info.name] = prop_array["data"]
                            else:
                                np_cell_data[info.name] = np.concatenate(
                                    (np_cell_data[info.name], prop_array["data"])
                                )
                            # if np_cell_data.get(info.name) is None:
                                # np_cell_data[info.name] = []
                            # np_cell_data[info.name].append(prop_array["data"])

                        else:
                            print(
                                f"Unknown data type {prop_array['type']} for property {info.name}"
                            )
                    else:
                        print(f"Property {info.name} is not on grid.")
                else:
                    print(f"Property {info.name} not found for timestep {tstep}.")

            # timestep_filename = f"{filename}_t{tstep}.vtkhdf"
            timestep_filename = f"{filename}_timeseries.vtkhdf"
            output_filepath = os.path.join(output_folder, timestep_filename)
            
            # get cell data info
            cell_data = grid.GetCellData()
            n_arrays = cell_data.GetNumberOfArrays()
            for i in range(n_arrays):
                array = cell_data.GetArray(i)
                print(f"Array {i}: {array.GetName()}")
                print(f"Array {i}: {array.GetNumberOfTuples()}")
                print(f"Array {i}: {array.GetNumberOfComponents()}")
                print(f"Array {i}: {array.GetDataType()}")

        # Add cell data
        # print(np_cell_data)
        for key, value in np_cell_data.items():
            data = numpy_support.numpy_to_vtk(value, deep=True)
            data.SetName(key)
            # grid.GetCellData().AddArray(data)

        timesteps.SetInputData(grid)

        print("Writing vtkhdf file")
        # writer = vtk.vtkHDFWriter()
        # writer.SetFileName(output_filepath)
        # writer.SetWriteAllTimeSteps(True)

        # writer.SetInputConnection(timesteps.GetOutputPort())
        # writer.Write()

        # Add cell data
        with h5.File(output_filepath, "w") as f:
            # write hdf file manually
            root = f.create_group("VTKHDF")
            typeAsString = "UnstructuredGrid".encode("ASCII")
            root.attrs.create("Type", typeAsString, dtype=h5.string_dtype("ASCII", len(typeAsString)))
            root.attrs["Version"] = (2, 1)

            root.create_dataset("Points", (0, 3), maxshape=(None, 3), dtype=np.float64)
            root.create_dataset("Connectivity", (0, ), maxshape=(None, ), dtype=np.int32)
            root.create_dataset("Offsets", (0, ), maxshape=(None, ), dtype=np.int32)
            root.create_dataset("Types", (0, ), maxshape=(None, ), dtype=np.ubyte)

            steps = root.create_group("Steps")
            steps.attrs["NSteps"] = nsteps

            # Cell data
            cell_data_group = root.create_group("CellData")
            cell_data_offsets = steps.create_group("CellDataOffsets")

            # utility function
            def append_dataset(dset, array):
                np_array = numpy_support.vtk_to_numpy(array)
                originalLength = dset.shape[0]
                dset.resize(originalLength + np_array.shape[0], axis=0)
                dset[originalLength:] = np_array

            # Write real data
            points = grid.GetPoints().GetData()
            append_dataset(root["Points"], points)
            root["NumberOfPoints"] = np.ones(nsteps, dtype=np.int32) * points.GetNumberOfTuples()

            connectivity = grid.GetCells().GetConnectivityArray()
            append_dataset(root["Connectivity"], connectivity)
            root["NumberOfCells"] = np.ones(nsteps, dtype=np.int32) * grid.GetNumberOfCells()
            root["NumberOfConnectivityIds"] = np.ones(nsteps, dtype=np.int32) * grid.GetCells().GetNumberOfConnectivityIds()
            append_dataset(root["Offsets"],grid.GetCells().GetOffsetsArray())
            append_dataset(root["Types"], grid.GetCellTypesArray())

            # Write steps
            steps["CellOffsets"] = np.zeros(nsteps, dtype=np.int32)
            steps["ConnectivityIdOffsets"] = np.zeros(nsteps, dtype=np.int32)
            steps["PartOffsets"] = np.zeros(nsteps, dtype=np.int32)
            steps["PointOffsets"] = np.zeros(nsteps, dtype=np.int32)

            # Write cell data
            for key, value in np_cell_data.items():
                cell_data_group.create_dataset(key, data=value, maxshape=(None, ), dtype=np.float64)
                cell_data_offsets.create_dataset(key, data=np.arange(0, nsteps * grid.GetNumberOfCells(), grid.GetNumberOfCells()), dtype=np.int64)

            steps["Values"] = np.array(time_values, dtype=np.float64)

            print(f"Successfully wrote vtkhdf file to {output_filepath}")



    def write_timeseries_complete(self, output_folder, filename, write_static=True, write_timesteps=True):
        """
        Write vtkhdf files for all timesteps. Support writing static and transient data.

        :param output_folder: Directory path for output files.
        :param filename: Base filename for output files.
        :param write_static: Write static data to vtkhdf file.
        :param write_timesteps: Write transient data to vtkhdf file.
        """

        os.makedirs(output_folder, exist_ok=True)
        timeseries_filename = f"{filename}_timeseries.vtkhdf"
        output_filepath = os.path.join(output_folder, timeseries_filename)

        # HDF chunk size, can be adjusted. Default: 25000 in paraview
        chunk_size = 262144
        # cell_chunk_size = 242430 # 242430 cells per chunk
        # point_chunk_size = 25000 # 25000 points per chunk

        print("Write vtkhdf file!")
        with h5.File(output_filepath, "w") as f:
            # Write hdf file header manually
            root = f.create_group("VTKHDF")
            typeAsString = "UnstructuredGrid".encode("ASCII")
            root.attrs.create("Type", typeAsString, dtype=h5.string_dtype("ASCII", len(typeAsString)))
            root.attrs["Version"] = (2, 1) # Same version as paraview used

            # Add static grid data

            # Process points
            print("Add points to vtkhdf file")
            GRID_VTX_array = {}
            if self.reader.get_static_data("GRID_VTX", GRID_VTX_array):
                points_array = GRID_VTX_array["data"].reshape(-1, 3)
                num_points = points_array.shape[0]
                root.create_dataset("Points", data=points_array, maxshape=(None, 3), dtype=np.float64, chunks=(chunk_size, 3))

            GRID_VTX_array.clear()

            # Process cells, only hexahedron
            # for final case, we can use int32 except for the cell data offset
            print("Add cell connectivity to vtkhdf file")
            NODE_LIST_array = {}
            if self.reader.get_static_data("NODE_LIST", NODE_LIST_array):
                # issue: use double memory for connectivity because of the np read-only array
                connectivity = NODE_LIST_array["data"] - 1 # Convert to zero-based indices
                num_connectivity = len(connectivity)
                num_cells = num_connectivity // 8
                root.create_dataset("Connectivity", data=connectivity, maxshape=(None, ), dtype=np.int32, chunks=chunk_size)

            NODE_LIST_array.clear()

            print("Add cell offsets to vtkhdf file")
            offsets = np.arange(num_cells + 1, dtype=np.int32) * 8 # 8 nodes per cell
            root.create_dataset("Offsets", data=offsets, maxshape=(None, ), dtype=np.int32, chunks=chunk_size)
            offsets = None

            print("Add cell types to vtkhdf file")
            types = np.full(num_cells, 12, dtype=np.ubyte) # vtk_hexahedron = 12
            root.create_dataset("Types", data=types, maxshape=(None, ), dtype=np.ubyte, chunks=chunk_size)
            types = None

            # Timestep properties
            nsteps = self.reader.get_time_steps()
            print(f"Number of timesteps: {nsteps}")
            time_values = []
            if (write_timesteps):
                # Write number properties
                root["NumberOfPoints"] = np.full(nsteps, num_points, dtype=np.int32)
                root["NumberOfConnectivityIds"] = np.full(nsteps, num_connectivity, dtype=np.int32)
                root["NumberOfCells"] = np.full(nsteps, num_cells, dtype=np.int32)

                # Write timesteps data
                steps = root.create_group("Steps")
                steps.attrs["NSteps"] = nsteps
                # grid data not changed
                steps.create_dataset("CellOffsets", data=np.zeros(nsteps, dtype=np.int64), dtype=np.int64)
                steps.create_dataset("ConnectivityIdOffsets", data=np.zeros(nsteps, dtype=np.int64), dtype=np.int64)
                steps.create_dataset("PartOffsets", data=np.zeros(nsteps, dtype=np.int64), dtype=np.int64)
                steps.create_dataset("PointOffsets", data=np.zeros(nsteps, dtype=np.int64), dtype=np.int64)

            else:
                # Not need to write timesteps data
                nsteps = 1
                root["NumberOfPoints"] = np.full(nsteps, num_points, dtype=np.int32)
                root["NumberOfConnectivityIds"] = np.full(nsteps, num_connectivity, dtype=np.int32)
                root["NumberOfCells"] = np.full(nsteps, num_cells, dtype=np.int32)

            # Process cell data
            cell_data_group = root.create_group("CellData")           
            if (write_timesteps):                
                cell_data_offsets = steps.create_group("CellDataOffsets")

            # Write static cell data
            if (write_static):
                for static_infos in self.reader.stat_props:
                    prop_array = {}
                    if self.reader.get_static_data(static_infos.name, prop_array):
                        if prop_array["on_grid"]:
                            # make sure the data is for all cells
                            if prop_array["type"] not in ["d", "i"]:
                                print(
                                    f"Unknown data type {prop_array['type']} for property {static_infos.name}"
                                )
                            elif prop_array["data"].shape[0] != num_cells:
                                print(
                                    f"Data shape {prop_array['data'].shape[0]} does not match number of cells {num_cells} for property {static_infos.name}"
                                )
                            else:
                                cell_data_group.create_dataset(static_infos.name, data=prop_array["data"], maxshape=(None, ), chunks=chunk_size) # no dtype
                                if (write_timesteps):
                                    cell_data_offsets.create_dataset(static_infos.name, data=np.zeros(nsteps, dtype=np.int64), dtype=np.int64) # no offset for static data
                                print(f"Adding property {static_infos.name} to grid")
                        else:
                            print(f"Property {static_infos.name} is not on grid.")
                    else:
                        print(f"Property {static_infos.name} not found.")

            # Write transient(timesteps) cell data
            if (write_timesteps):
                dym_datasets = {}

                for ind, dynamic_infos in enumerate(self.reader.dym_props):
                    tstep = np.float64(dynamic_infos[0].time)
                    if ind >= len(self.reader.dym_props):
                        print(f"Time step {tstep} does not exist. Skipping.")
                        continue
                    time_values.append(tstep)

                    for info in dynamic_infos:
                        prop_array = {}
                        if self.reader.get_dynamic_data(ind, info.name, prop_array):
                            if prop_array["on_grid"]:
                                if prop_array["type"] in ["d", "i"]:
                                    print(f"Adding property {info.name} for timestep {tstep} to grid")
                                    data = prop_array["data"]

                                    if info.name not in dym_datasets:
                                        # create expand dataset
                                        dset = cell_data_group.create_dataset(
                                            info.name,
                                            shape=(0,),
                                            maxshape=(None,),
                                            chunks=(chunk_size,),
                                            dtype=data.dtype
                                        )
                                        dym_datasets[info.name] = dset
                                    else:
                                        dset = dym_datasets[info.name]

                                    current_size = dset.shape[0]
                                    dset.resize(current_size + data.shape[0], axis=0)
                                    dset[current_size:current_size + data.shape[0]] = data
                                else:
                                    print(f"Unknown data type {prop_array['type']} for property {info.name}")
                            else:
                                print(f"Property {info.name} is not on grid.")
                        else:
                            print(f"Property {info.name} not found for timestep {tstep}.")

                for key, value in dym_datasets.items():
                    # cell_data_group.create_dataset(key, data=value, maxshape=(None, ), chunks=chunk_size) # no dtype
                    cell_data_offsets.create_dataset(key, data=np.arange(nsteps) * num_cells, dtype=np.int64) # can not use np.int32

                # Write time values
                steps.create_dataset("Values", data=np.array(time_values, dtype=np.float64), dtype=np.float64)

            print(f"Successfully wrote vtkhdf file to {output_filepath}")

if __name__ == "__main__":
    print("This is a writer class file. Please run main.py for execution")
